#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
模型下载脚本
用于下载embedding模型到本地models目录
"""

import os
import sys
from pathlib import Path
from sentence_transformers import SentenceTransformer
import yaml


def load_config():
    """加载配置文件"""
    config_path = Path("configs/llm_config.yaml")
    if not config_path.exists():
        print(f"配置文件不存在: {config_path}")
        return None
    
    with open(config_path, 'r', encoding='utf-8') as f:
        return yaml.safe_load(f)


def download_embedding_model(model_name: str, target_path: str):
    """
    下载embedding模型到指定路径
    
    Args:
        model_name: HuggingFace模型名称
        target_path: 目标下载路径
    """
    # 确保使用相对路径，基于当前工作目录
    target_path = Path(target_path)
    if not target_path.is_absolute():
        target_path = Path.cwd() / target_path
    
    # 确保目标目录存在
    target_path.mkdir(parents=True, exist_ok=True)
    
    print(f"开始下载模型: {model_name}")
    print(f"目标路径: {target_path}")
    
    try:
        # 先下载模型到临时位置
        print("正在从HuggingFace下载模型...")
        model = SentenceTransformer(model_name)
        
        # 将模型保存到指定路径
        print(f"正在保存模型到: {target_path}")
        model.save(str(target_path))
        
        print(f"✅ 模型下载成功！")
        print(f"模型已保存到: {target_path}")
        
        # 验证模型是否可以正常加载
        print("正在验证模型...")
        test_model = SentenceTransformer(str(target_path))
        test_embedding = test_model.encode("测试文本")
        print(f"✅ 模型验证成功！embedding维度: {len(test_embedding)}")
        
        return True
        
    except Exception as e:
        print(f"❌ 模型下载失败: {e}")
        return False


def download_tts_models():
    """下载TTS相关模型（如果需要的话）"""
    print("TTS模型需要手动下载，请根据配置中的路径放置模型文件：")
    print("- GPT_MODEL_PATH: models/雷电将军.ckpt")
    print("- SOVITS_MODEL_PATH: models/雷电将军.pth") 
    print("- REF_AUDIO_PATH: models/雷电将军/我此番也是受神子之邀，体验一下市井游乐的氛围，和各位并无二致。.wav")


def main():
    """主函数"""
    print("🚀 开始下载模型...")
    
    # 加载配置
    config = load_config()
    if not config:
        return
    
    embedding_config = config.get('embedding_config', {})
    model_name = embedding_config.get('model_name', 'BAAI/bge-small-zh-v1.5')
    model_path = embedding_config.get('model_path', 'models/embedding')
    
    print(f"配置的模型名称: {model_name}")
    print(f"配置的模型路径: {model_path}")
    
    # 检查是否已存在本地模型
    model_path_obj = Path(model_path)
    if not model_path_obj.is_absolute():
        model_path_obj = Path.cwd() / model_path_obj
    
    if model_path_obj.exists() and (model_path_obj / "config.json").exists():
        print(f"⚠️  本地模型已存在: {model_path_obj}")
        response = input("是否要重新下载？(y/N): ").strip().lower()
        if response != 'y':
            print("跳过下载")
            return
    
    # 下载embedding模型
    success = download_embedding_model(model_name, model_path)
    
    if success:
        print("\n🎉 所有模型下载完成！")
        print("现在您的代码将使用本地模型，无需联网下载。")
    else:
        print("\n❌ 模型下载失败，请检查网络连接和配置。")
    
    # 提示TTS模型
    print("\n" + "="*50)
    download_tts_models()


if __name__ == "__main__":
    main()
